import argparse

import torch

from train import prepare


def main():
    parser = argparse.ArgumentParser("train model")
    parser.add_argument('--task', type=str, help='task type',
                        default='mnist', choices=['mnist', 'cifar10'])
    parser.add_argument('--dataset-root', type=str)
    parser.add_argument('--ckpt-path', type=str,
                        help='checkpoint path')
    args = parser.parse_args()

    model, _, testloader = prepare(args.task, args.dataset_root)
    model.eval()
    model.load_state_dict(torch.load(args.ckpt_path))

    correct = 0
    total = 0
    with torch.no_grad():
        for i, data in enumerate(testloader):
            inputs, labels = data
            inputs, labels = inputs.to("cuda"), labels.to("cuda")
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'accuracy: {correct / total}')


if __name__ == "__main__":
    main()
